#! /usr/bin/env python
# coding=utf-8
#======================================================================
# SecureGAppProxy is a security-strengthened version of GAppProxy.
# http://secure-gappproxy.googlecode.com                               
# This file is a part of SecureGAppProxy.                              
# Copyright (C) 2011  nleven <www.nleven.com i@nleven.com>             
#                                                                      
# This program is free software: you can redistribute it and/or modify 
# it under the terms of the GNU General Public License as published by 
# the Free Software Foundation, either version 3 of the License, or    
# (at your option) any later version.                                  
#                                                                      
# This program is distributed in the hope that it will be useful,      
# but WITHOUT ANY WARRANTY; without even the implied warranty of       
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the        
# GNU General Public License for more details.                         
#                                                                      
# You should have received a copy of the GNU General Public License    
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#                                                                      
# ACKNOWLEDGEMENT                                                      
# SecureGAppProxy is a based on the work of GAppProxy                  
# <http://gappproxy.googlecode.com> by Du XiaoGang <dugang@188.com>
#======================================================================

import pkcs5
import base64
import getpass

import os, sys
import re

from SuffixTree.SubstringDict import getAllIndices
from SuffixTree import SuffixTree

import marshal


pwddict = None
bktree = None
suffixtree = None


def edit_dist(a, b):
    d = list(range(len(a)+1))
    new_d = [0] * (len(a)+1)
    for i in xrange(1, len(b)+1):
        new_d[0] = i
        for j in xrange(1, len(a)+1):
            if j <= len(a) and i <= len(b) and a[j-1]==b[i-1]:
                new_d[j] = d[j-1]
            else:
                new_d[j] = min(d[j]+1, new_d[j-1]+1, d[j-1]+1)
        new_d, d = d, new_d
    return d[len(a)]

def load_dictionary():
    global pwddict, bktree, suffixtree
    dict_path = os.path.join(os.path.dirname(sys.argv[0]),'pwddict')
    f = open(dict_path)
    pwddict = f.read().split('\n')
    f.close()
    
    bktree_path=os.path.join(os.path.dirname(sys.argv[0]),'pwddict_bktree')
    f = open(bktree_path, 'rb')
    bktree = marshal.load(f)
    f.close()

    suffixtree = SuffixTree()
    for i, pwd in enumerate(pwddict):
        suffixtree.add(pwd.lower(), i)

def find_closest_edit_dist(password, min_dist):
    closest_dist, closest_pwd= len(password), ''
    
    node_list = [bktree]
    while len(node_list) != 0:
        node = node_list.pop(0)
        node_pwd = pwddict[node[0]]
        
        dist = edit_dist(node_pwd, password)
        if dist <= closest_dist:
            closest_dist, closest_pwd = dist, node_pwd
        
        for i in range(max(1,dist-min_dist), dist+min_dist+1):
            if i in node:
                node_list.append(node[i])
    
    return closest_dist, closest_pwd

def find_closest_lcs(password):
    max_num_matched, closest_node = 0, None
    for i in range(len(password)):
        num_matched, node, _ = suffixtree.match(password[i:])
        if max_num_matched < num_matched:
            max_num_matched, closest_node = num_matched, node

    closest_pwd = ''
    if closest_node:
        pwd_candidate = getAllIndices(closest_node)
        closest_pwd = min( map(lambda i:pwddict[i], pwd_candidate), key=lambda p:len(p))            

    return max_num_matched, closest_pwd 


def estimate_complexity(password):
    import math
    
    m = re.match('^(?P<notnum>[^0-9]+)(?P<num>[0-9]+)$', password)
    if m:
        return -math.log(0.3,2) + estimate_complexity(m.group('notnum')) + estimate_complexity(m.group('num'))
    
    return math.log(estimate_charset_size(password) ** len(password), 2)

def estimate_charset_size(password):
    pwd_space = 0
    if re.search('[a-z]', password):
        pwd_space += 26
    if re.search('[A-Z]', password):
        pwd_space += 26
    if re.search('[0-9]', password):
        pwd_space += 10
    #common english punctuation
    if re.search('[\.\, ]', password):
        pwd_space += 3
    #symbols on the first row
    if re.search('[!@£#\$%\^&\*\(\)\-_=\+]', password):
        pwd_space += 15
    #other symbols
    if re.search("""[\?\/><`~\\|"';:\]\}\[\{\s]""", password):
        pwd_space += 17
    return pwd_space

def is_keyboard_related(password):
    key =       ['`1234567890-=', '\tqwertyuiop[]\\', 'asdfghjkl;\'', 'zxcvbnm,./']
    key_shift = ['~!@#$%^&*()_+', '\tQWERTYUIOP{}|',  'ASDFGHJKL:"',  'ZXCVBNM<>?']
    row_offset = [0, 0, 1, 1, 1]
    def char_to_pos(c):
        for i in range(len(key)):
            if c in key[i]:
                key_t = key[i]
            elif c in key_shift[i]:
                key_t = key_shift[i]
            else:
                continue
            return (i, row_offset[i]+key_t.find(c))
        return (0, 0)
    key_pos = map(char_to_pos, password)

    def is_neighbour(pos1, pos2):
        x1, y1 = pos1
        x2, y2 = pos2
        return (abs(x1-x2)<=1 and abs(y1-y2)<=1)

    neighbour_count = 0
    for i in range(len(key_pos[:-1])):
        if is_neighbour(key_pos[i], key_pos[i+1]):
            neighbour_count += 1
    return len(password)-neighbour_count <= 3
        
    

def filter_repitition(password, least_len=3):
    if least_len > len(password)/2:
        return len(password)
    mask = [True]*len(password)
    for i in xrange(len(password)-least_len):
        sub = password[i:i+least_len]
        if password.count(sub) > 1:
            start = password.find(sub)
            while start!=-1:
                for j in xrange(start, start+least_len):
                    mask[j] = False
                start = password.find(sub, start+1)
    minlen = 0
    for j in xrange(len(password)):
        if mask[j]:
            minlen+=1
    return min(minlen, filter_repitition(password, least_len+1))

def get_password_hash(server):
    load_dictionary()
    
    pwd = ''
    while True:
        pwd = getpass.getpass('Password for SecureGAppProxy: ')
        
        if len(pwd) < 7:
            print "Your password is too short. Passwords must be at least 7 in length. Please try again."
            continue
        
        if re.match("^[0-9]+$", pwd):
            print "Your password only consists of numbers. Consider adding symbols and letters."
            continue
        
        if estimate_complexity(pwd) < 40:
            print "Your password is too simple. Consider adding numbers, symbols and letters in mixed case."
            continue
        
        if filter_repitition(pwd) <= 4:
            print "Your password contains too many repititions. Please try another one."
            continue
        
        length, closest = find_closest_lcs(pwd.lower())
        if len(pwd)-length <=3 or (len(pwd) >= 15 and length >= len(pwd)/2):
            print "Your password looks similar to a known weak password. Please try another one."
            if raw_input("See the weak password?(y/n)") in ('Y', 'y'):
                print closest
            continue

        if is_keyboard_related(pwd):
            print "Your password appears to be a common sequence on the keyboard. Please try another one."
            continue

        dist, closest = find_closest_edit_dist(pwd.lower(), 3)
        if dist <= 3:
            print "Your password looks similar to a known weak password. Please try another one."
            if raw_input("See the weak password?(y/n)") in ('Y', 'y'):
                 print closest
            continue
            
        pwd_confirm = getpass.getpass('Password accepted, please confirm password: ')
        if pwd != pwd_confirm:
            print "Two passwords aren't consistent, please try again."
            continue

        break
        
    print "Computing hash...\n"
    storage = pkcs5.PBKDF1(pwd, server.lower())
        
    return base64.urlsafe_b64encode(storage)

if __name__ == '__main__':
    print get_password_hash('test')
